package cn.felord.springbootwebsocket.config;

import com.sun.security.auth.UserPrincipal;
import org.springframework.messaging.Message;
import org.springframework.messaging.MessageChannel;
import org.springframework.messaging.simp.stomp.StompCommand;
import org.springframework.messaging.simp.stomp.StompHeaderAccessor;
import org.springframework.messaging.support.ChannelInterceptor;
import org.springframework.messaging.support.MessageHeaderAccessor;
import org.springframework.util.StringUtils;

import java.security.Principal;
import java.util.Objects;

/**
 * @author dax
 * @since 2019/4/14 0:02
 */
public class HeaderChannelInterceptor implements ChannelInterceptor {
    @Override
    public Message<?> preSend(Message<?> message, MessageChannel channel) {
        StompHeaderAccessor accessor = MessageHeaderAccessor.getAccessor(message, StompHeaderAccessor.class);

        if (StompCommand.CONNECT.equals(Objects.requireNonNull(accessor).getCommand())) {
            String username = accessor.getFirstNativeHeader("username");
            if (StringUtils.hasLength(username)) {
                // 绑定user
                Principal principal = new UserPrincipal(username);
                accessor.setUser(principal);
            }

        }
        return message;
    }
}
